import torch
import numpy as np
from PIL import Image
import base64
import io
import json
import os
from volcenginesdkarkruntime import Ark

# 获取当前文件所在目录
current_dir = os.path.dirname(__file__)
# 构建 config.json 的完整路径
config_path = os.path.join(current_dir, 'config.json')

# 加载配置文件
config = {}
if os.path.exists(config_path):
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
    except Exception as e:
        print(f"Kontext Analyze: Error loading config.json: {e}")

class KontextDuoImageAnalyzer:
    @classmethod
    def INPUT_TYPES(s):
        """
        定义节点的输入类型和控件。
        现在会从 config.json 读取默认值。
        """
        return {
            "required": {
                "image_a": ("IMAGE",),
                "image_b": ("IMAGE",),
                "api_key": ("STRING", {
                    "multiline": False,
                    "default": config.get("api_key", "YOUR_ARK_API_KEY_HERE")
                }),
                "model_id": ("STRING", {
                    "multiline": False,
                    "default": config.get("model_id", "ep-20250705181415-gkgwc")
                }),
                "base_url": ("STRING", {
                    "multiline": False,
                    "default": config.get("base_url", "https://ark.cn-beijing.volces.com/api/v3")
                }),
                "prompt": ("STRING", {
                    "multiline": True,
                    "default": "请对比分析这两张图片，总结它们之间的核心差异和共同点。"
                }),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("analysis_text",)
    FUNCTION = "analyze"
    CATEGORY = "Kontext"

    def tensor_to_pil(self, tensor: torch.Tensor) -> Image.Image:
        """
        将 ComfyUI 的图像张量 (Tensor) 转换为 PIL Image 对象。
        明确处理批处理，只取第一张图。
        """
        # 取出批次中的第一张图，并从 [0,1] 范围转换为 [0,255] 范围
        image_np = tensor[0].cpu().numpy() * 255.0
        image_np = np.clip(image_np, 0, 255).astype(np.uint8)

        # 从 Numpy 数组创建 PIL Image
        return Image.fromarray(image_np)

    def pil_to_base64(self, pil_image: Image.Image) -> str:
        """
        将 PIL Image 对象编码为 Base64 字符串。
        """
        buffered = io.BytesIO()
        pil_image.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode('utf-8')

    def analyze(self, image_a, image_b, api_key, model_id, base_url, prompt):
        """
        核心分析函数 - 使用新的 Ark SDK 并增加健壮性检查。
        """
        if not api_key or "YOUR_ARK_API_KEY_HERE" in api_key:
            return ("错误：请输入有效的 API Key，您可以在节点的输入框中或在 config.json 文件中提供。",)
        
        if not base_url:
            return ("错误：请输入有效的 Base URL。",)

        try:
            # 1. 将输入的张量转换为 PIL 图像
            pil_a = self.tensor_to_pil(image_a)
            pil_b = self.tensor_to_pil(image_b)

            # 2. 将 PIL 图像编码为 Base64
            base64_a = self.pil_to_base64(pil_a)
            base64_b = self.pil_to_base64(pil_b)

            # 3. 初始化 Ark 客户端
            client = Ark(api_key=api_key, base_url=base_url)

            # 4. 发送请求
            response = client.chat.completions.create(
                model=model_id,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/png;base64,{base64_a}"}
                            },
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/png;base64,{base64_b}"}
                            },
                        ],
                    }
                ],
            )

            # 5. 安全地提取结果
            if response.choices and len(response.choices) > 0:
                result_text = response.choices[0].message.content
                return (result_text,)
            else:
                return ("错误：API 返回了空的响应。",)

        except Exception as e:
            # 返回详细的错误信息
            error_message = f"Kontext Analyze: 分析时出现错误: {str(e)}"
            print(error_message) # 在控制台也打印错误，方便调试
            return (error_message,)

# ComfyUI 加载节点所必需的字典
NODE_CLASS_MAPPINGS = {
    "KontextDuoImageAnalyzer": KontextDuoImageAnalyzer
}

# 方便在菜单中显示的名称
NODE_DISPLAY_NAME_MAPPINGS = {
    "KontextDuoImageAnalyzer": "Kontext Duo Image Analyzer"
} 